package ml;

import weka.core.*;
import weka.classifiers.Classifier;

import java.io.FileInputStream;
import java.io.FileReader;
import java.io.ObjectInputStream;
import java.io.OutputStream;
import java.io.PrintStream;
import java.util.List;

/**
 * Classifier for weka 3.7. It uses sparse data.
 */
public class WekaGlue extends MLGlue {

	private Classifier cls;
	private Instances dataset;

	static {
		PrintStream error = System.err;
		PrintStream out = System.out;
		System.setErr(new PrintStream(new OutputStream() {
			public void write(int b) {
			}
		}));
		System.setOut(new PrintStream(new OutputStream() {
			public void write(int b) {
			}
		}));
		weka.core.WekaPackageManager.loadPackages(false);
		System.setErr(error);
		System.setOut(out);
	}

	/**
	 * @param modelName
	 *            Trained model
	 * @param datasetName
	 *            Arff file
	 */
	public WekaGlue(String modelName, String datasetName) throws Exception {
		ObjectInputStream ois = new ObjectInputStream(new FileInputStream(
				modelName));
		cls = (Classifier) ois.readObject();
		ois.close();

		dataset = new Instances(new FileReader(datasetName));
		dataset.setClassIndex(dataset.numAttributes() - 1);
	}

	public int getNAttributes() {
		return dataset.numAttributes() - 1;
	}

	@Override
	public String classify(List<Data> data) throws Exception {
		Instance instance = new SparseInstance(data.size());
		instance.setDataset(dataset);

		for (int i = 0; i < data.size(); i++) {
			try {
				if (data.get(i).getColumn() < 0)
					continue;
				instance.setValue(data.get(i).getColumn(), data.get(i)
						.getValue());
			} catch (IllegalArgumentException e) {
				System.err.println("Couldn't use argument "
						+ data.get(i).getColumn() + " to classifier: "
						+ data.get(i).getValue());
				System.err.println("Probably not defined in ARFF header.");
				e.printStackTrace();
				System.err.flush();
				return null;
			}
		}
		double d = cls.classifyInstance(instance);
		return dataset.classAttribute().value((int) d);
	}
}
